feat: add Bedrock InvokeModelWithResponseStream instrumentation #2845
Conversation
yuliia-fryshko
force-pushed
the
bedrock-invoke-model-stream
branch
2 times, most recently
from
09f9777 to
eeb1e84
Compare
| @@ -102,6 +102,13 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { | |||
| return this.requestPreSpanHookConverse(request, config, diag, true); | |||
| case 'InvokeModel': | |||
| return this.requestPreSpanHookInvokeModel(request, config, diag); | |||
| case 'InvokeModelWithResponseStream': | |||
| return this.requestPreSpanHookInvokeModelWithResponseStream( | |||
There was a problem hiding this comment.
Is there a reason to not re-use requestPreSpanHookInvokeModel: add a 4th isStream argument and pass in false for 'InvokeModel', true for 'InvokeModelWithResponseStream', and then make the minor update to the implementation? This is how it was done for 'Converse' and 'ConverseStream'.
It looks to me like the requestPreSpanHookInvokeModel and requestPreSpanHookInvokeModelWithResponseStream functions are almost identical ... except that the latter doesn't have blocks for 'meta.llama', 'cohere.*', and 'mistral'.
There was a problem hiding this comment.
Thanks for the suggestion, @trentm ! You are absolutely right! I've updated the code to consolidate requestPreSpanHookInvokeModel and requestPreSpanHookInvokeModelWithResponseStream into a single method using isStream parameter as you suggested.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2845 +/- ##
==========================================
+ Coverage 89.68% 89.73% +0.04%
==========================================
Files 187 187
Lines 9059 9118 +59
Branches 1858 1876 +18
==========================================
+ Hits 8125 8182 +57
- Misses 934 936 +2
🚀 New features to boost your workflow:
|
yuliia-fryshko
force-pushed
the
bedrock-invoke-model-stream
branch
from
eeb1e84 to
0d197f0
Compare
| ): Promise<any> { | ||
| const stream = response.data?.body; | ||
| const modelId = response.request.commandInput?.modelId; | ||
| if (!stream || !span.isRecording()) return; |
There was a problem hiding this comment.
| if (!stream || !span.isRecording()) return; | |
| if (!stream) return; |
!span.isRecording() is already checked before responseHookInvokeModelWithResponseStream() is called
| return response.data; | ||
|
|
||
| // Tap into the stream at the chunk level without modifying the chunk itself. | ||
| function instrumentAsyncIterable<T>( |
There was a problem hiding this comment.
Wouldn't it be more efficient to declare these following functions outside of this member function scope?
Which problem is this PR solving?
Adds instrumentation of the InvokeModelWithResponseStreamCommand in the AWS Bedrock SDK.
Short description of the changes
instrumentAsyncIterableis used to inspect streamed chunks in real time and extract relevant telemetry.